package com.xialuo.study.std;

import com.xialuo.study.state.ReactState;
import lombok.extern.slf4j.Slf4j;
import org.bsc.langgraph4j.serializer.std.ObjectStreamStateSerializer;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.ToolResponseMessage;
import org.springframework.ai.chat.messages.UserMessage;

@Slf4j
public class StateSerializer extends ObjectStreamStateSerializer<ReactState> {

    public StateSerializer() {
        super(ReactState::new);
        mapper().register(SystemMessage.class, new SystemMessageSerializer());
        mapper().register(UserMessage.class, new UserMessageSerializer());
        mapper().register(AssistantMessage.ToolCall.class, new ToolCallSerializer());
        mapper().register(ToolResponseMessage.ToolResponse.class, new ToolResponseSerializer());
        mapper().register(AssistantMessage.class, new AssistantMessageSerializer());
        mapper().register(ToolResponseMessage.class, new ToolResponseMessageSerializer());
    }
}
